[https://nvbugs/6501376][fix] Test-only fix — drop the if hidden_size % 2 != 0: with pytest.raises(...)…#16844
[https://nvbugs/6501376][fix] Test-only fix — drop the if hidden_size % 2 != 0: with pytest.raises(...)…#16844trtllm-agent wants to merge 4 commits into
if hidden_size % 2 != 0: with pytest.raises(...)…#16844Conversation
…mn linear PR NVIDIA#14875 added VANILLA-mode uneven-TP support in tensorrt_llm/_torch/modules/linear.py via _auto_tp_sharding, so test_row_linear[2-unbalanced] (hidden=15) and test_column_linear [2-unbalanced] (hidden=127) no longer hit the divisibility assert. The pre-existing tests still expected an AssertionError on the odd branch, producing "Failed: DID NOT RAISE" in pre-merge CI. Unconditionally assert all ranks return True. Signed-off-by: handongl <handongl@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Walkthrough
ChangesUneven linear sharding
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant LinearInit
participant AutoTPSharding
participant LinearForward
participant Allgather
LinearInit->>AutoTPSharding: Query per-rank shard widths
AutoTPSharding-->>LinearInit: Return shard boundaries
LinearForward->>Allgather: Pass gather_output_sizes
Allgather-->>LinearForward: Gather uneven column outputs
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
|
/bot run |
|
PR_Github #61699 [ run ] triggered by Bot. Commit: |
|
PR_Github #61699 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61711 [ run ] triggered by Bot. Commit: |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot kill |
…allgather Signed-off-by: Lori Ren <lorir@nvidia.com>
|
PR_Github #61714 [ kill ] triggered by Bot. Commit: |
|
PR_Github #61711 [ run ] completed with state |
|
PR_Github #61714 [ kill ] completed with state |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/linear.py (1)
3415-3415: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the expanded helper signature.
Add precise parameter and return annotations to
_auto_tp_sharding, including the new optionalrankargument.As per coding guidelines, every function must be annotated.
Proposed signature
- def _auto_tp_sharding(self, features, quant_config, rank=None): + def _auto_tp_sharding( + self, + features: int, + quant_config: Optional[QuantConfig], + rank: Optional[int] = None, + ) -> tuple[int, int]:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/linear.py` at line 3415, Annotate the _auto_tp_sharding method signature with precise types for features, quant_config, and the optional rank parameter, and add its return annotation. Use the existing project type conventions and ensure rank’s optionality is explicitly represented.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Around line 3353-3360: Extend the gather_output validation around the tuple
check to handle dictionary-based self.tp_sharding overrides as well. When
dictionary sharding produces uneven local widths, either derive a complete
per-rank sizes list and pass it to allgather via gather_output_sizes, or raise
ValueError for the unsupported case; do not leave gather_output_sizes unset and
allow equal-width assumptions.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Line 3415: Annotate the _auto_tp_sharding method signature with precise types
for features, quant_config, and the optional rank parameter, and add its return
annotation. Use the existing project type conventions and ensure rank’s
optionality is explicitly represented.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 248bddb6-20d7-4efa-8400-9eb5536a9bcb
📒 Files selected for processing (1)
tensorrt_llm/_torch/modules/linear.py
…sharding Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #61716 [ run ] triggered by Bot. Commit: |
|
PR_Github #61716 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61721 [ run ] triggered by Bot. Commit: |
|
PR_Github #61721 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61732 [ run ] triggered by Bot. Commit: |
|
PR_Github #61732 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61742 [ run ] triggered by Bot. Commit: |
Summary
pytest.raises(AssertionError)on the odd-hidden_sizebranch expected a divisibility assert that PR [TRTLLM-13117][feat] Implement Uneven TP Linear for VisualGen models #14875's_auto_tp_shardingno longer fires for VANILLA row/column Linear.if hidden_size % 2 != 0: with pytest.raises(...)branch in bothtest_row_linearandtest_column_linear; unconditionallyfor r in results: assert r is True.Test plan
Links
Dev Engineer Review
tensorrt_llm/_torch/modules/linear.pyto correctly handlegather_output=Truefor VANILLATensorParallelMode.COLUMNwhen_auto_tp_shardingproduces uneven per-rank shard widths:_tp_sharding_is_autotracking to distinguish whethertp_shardingcame from_auto_tp_shardingvsoverride_tp_sharding.gather_output_sizes(defaultNone) to store per-rankallgathersizes when uneven widths are detected._auto_tp_sharding(..., rank=...)to accept an optionalrankargument so shard boundaries can be queried for other ranks.forwardfor COLUMN +gather_output=True, passedsizes=self.gather_output_sizesinto the COLUMN allgather to avoid relying on shape inference.gather_outputwithoverride_tp_sharding(via aValueError), aligning with the referenced bug about this unsupported mode combination.tests/unittest/_torch/multi_gpu/test_linear.py:pytest.raises(AssertionError)expectations for “unbalanced/odd”hidden_sizein bothtest_column_linearandtest_row_linear.True, reflecting the post-[TRTLLM-13117][feat] Implement Uneven TP Linear for VisualGen models #14875 behavior where_auto_tp_shardingno longer triggers that assertion for odd hidden sizes.QA Engineer Review
tests/unittest/_torch/multi_gpu/test_linear.py):test_column_linear(removed conditional failure expectation; now assertsTruefor all ranks)test_row_linear(removed conditional failure expectation; now assertsTruefor all ranks)tests/integration/test_lists/waives.txt):unittest/_torch/multi_gpu/test_linear.py::test_column_linear[2-unbalanced]unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-unbalanced]